

package class10;
import java.util.HashMap;
public class Code_01_UnionFind {
    public static class Node<V> {
        V value;
        public Node(V v){
            value = v;
        }
    }

    public static class UnionSet<V>{
        public HashMap<V, Node<V>> nodes;
        public HashMap<Node<V>, Node<V>> parents;
        public HashMap<Node<V>, Integer> sizeMap;
    }

    public UnionSet(List<V> values){
        for (V cur : values) {
            Node<V> node = new Node(cur);
            nodes.put(cur, node);
            parents.put(node,node);
            sizeMap.put(node, 1);
        }
    }

    public Node<V> findFather(Node<V> cur){
        Stack<Node<V>> path = new Stack<Node<V>>();

        while(cur != parents.get(cur)){
            path.push(cur);
            cur = parents.get(cur);
        }

        //cur节点
        while(!path.isEmpty()){
            parents.put(path.pop(), cur);
        }
        return cur;
    }

    public boolean isSameSet(V a, V b){
        if(!nodes.containsKey(a) || !nodes.containsKey(b)){
            return false;
        }
        return findFather(nodes.get(a) == findFather(nodes.get(b)));
    }


    public void union(V a , V b){
        if (!nodes.containsKey(a)  || !nodes.containsKey(b)){
            return ;
        }

        Node<V> aHead = findFather.get(nodes.get(a));
        Node<V> bHead = findFather.get(nodes.get(b));

        if (aHead != bHead){
            int aSetSize = sizeMap.get(aHead);
            int bSetSize = sizeMap.get(bHead);
            Node<V> big = aSetSize >= bSetSize ? aHead : bHead;
            Node<V> small = big == aHead ? bHead : aHead;

            parents.put(small, big);
            sizeMap.put(big, aSetSize + bSetSize);
            sizeMap.remove(small);
        }
    }
}
